package cmd

import (
	"fmt"
	"math"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"mongo-stress/point"
	"mongo-stress/stress"
	"mongo-stress/write"

	"github.com/spf13/cobra"
)

var (
	// statsHost, statsDB                               string
	host, db, collection string
	// rp, precision, consistency 						string
	// username, password 								string
	createCommand, dump string
	seriesN             int
	// gzip                                    		 	int
	batchSize, pointsN, pps uint64
	runtime                 time.Duration
	tick                    time.Duration
	fast, quiet             bool
	strict, kapacitorMode   bool
	// recordStats                                      bool
	// tlsSkipVerify                                    bool
)

const (
	defaultSeriesKey string = "tags=sensor,region"
	defaultFieldStr  string = "value=0i"
)

var insertCmd = &cobra.Command{
	Use:   "insert DOC_TEMPLATE",
	Short: "Insert data into MongoDB",
	Long:  "Generate and insert JSON documents into MongoDB with configurable concurrency and throughput.",
	Run:   insertRun,
}

func insertRun(cmd *cobra.Command, args []string) {
	seriesKey := defaultSeriesKey
	fieldStr := defaultFieldStr
	if len(args) >= 1 {
		seriesKey = args[0]
		if !strings.Contains(seriesKey, ",") && !strings.Contains(seriesKey, "=") {
			fmt.Fprintln(os.Stderr, "First positional argument must be a series key, got: ", seriesKey)
			cmd.Usage()
			os.Exit(1)
			return
		}
	}
	if len(args) == 2 {
		fieldStr = args[1]
	}

	concurrency := pps / batchSize
	// PPS takes precedence over batchSize.
	// Adjust accordingly.
	if pps < batchSize {
		batchSize = pps
		concurrency = 1
	}

	if !quiet {
		fmt.Printf("Using point template: %s %s <timestamp>\n", seriesKey, fieldStr)
		fmt.Printf("Using batch size of %d line(s)\n", batchSize)
		fmt.Printf("Spreading writes across %d series\n", seriesN)
		if fast {
			fmt.Println("Output is unthrottled")
		} else {
			fmt.Printf("Throttling output to ~%d points/sec\n", pps)
		}
		fmt.Printf("Using %d concurrent writer(s)\n", concurrency)

		fmt.Printf("Running until ~%d points sent or until ~%v has elapsed\n", pointsN, runtime)
	}

	c := client()

	if !kapacitorMode {
		if err := c.Create(createCommand); err != nil {
			fmt.Fprintln(os.Stderr, "Failed to create database:", err.Error())
			fmt.Fprintln(os.Stderr, "Aborting.")
			os.Exit(1)
			return
		}
	}

	pts := point.NewPoints(seriesKey, fieldStr, seriesN)
	pointsPerSeries := pps / uint64(seriesN)
	interval := time.Second / time.Duration(pointsPerSeries)

	baseTimes := make([]time.Time, seriesN)
	counts := make([]uint64, seriesN)
	for i := 0; i < seriesN; i++ {
		baseTimes[i] = time.Now()
	}

	startSplit := 0
	if uint64(seriesN) < concurrency {
		concurrency = 1
	}
	inc := int(seriesN) / int(concurrency)
	endSplit := inc

	var wg sync.WaitGroup
	wg.Add(int(concurrency))

	var totalWritten uint64

	start := time.Now()
	for i := uint64(0); i < concurrency; i++ {

		go func(startSplit, endSplit int) {
			tick := time.Tick(tick)

			if fast {
				tick = time.Tick(time.Nanosecond)
			}

			cfg := stress.WriteConfig{
				BatchSize: batchSize,
				MaxPoints: pointsN / concurrency, // divide by concurreny
				Interval:  interval,
				BaseTimes: baseTimes,
				Counts:    counts,
				Offset:    startSplit,
				Deadline:  time.Now().Add(runtime),
				Tick:      tick,
			}

			// Ignore duration from a single call to Write.
			pointsWritten, _ := stress.Write(pts[startSplit:endSplit], c, cfg)
			atomic.AddUint64(&totalWritten, pointsWritten)

			wg.Done()
		}(startSplit, endSplit)

		startSplit = endSplit
		endSplit += inc
	}

	wg.Wait()
	totalTime := time.Since(start)
	if err := c.Close(); err != nil {
		fmt.Fprintf(os.Stderr, "Error closing client: %v\n", err.Error())
	}

	throughput := int(float64(totalWritten) / totalTime.Seconds())
	if quiet {
		fmt.Println(throughput)
	} else {
		fmt.Println("Write Throughput:", throughput)
		fmt.Println("Points Written:", totalWritten)
	}
}

func init() {
	RootCmd.AddCommand(insertCmd)
	// insertCmd.Flags().StringVarP(&statsHost, "stats-host", "", "http://localhost:8086", "Address of InfluxDB instance where runtime statistics will be recorded")
	// insertCmd.Flags().StringVarP(&statsDB, "stats-db", "", "stress_stats", "Database that statistics will be written to")
	// insertCmd.Flags().BoolVarP(&recordStats, "stats", "", false, "Record runtime statistics")
	insertCmd.Flags().StringVarP(&host, "host", "", "http://localhost:8086", "Address of InfluxDB instance")
	// insertCmd.Flags().StringVarP(&username, "user", "", "", "User to write data as")
	// insertCmd.Flags().StringVarP(&password, "pass", "", "", "Password for user")
	insertCmd.Flags().StringVarP(&db, "db", "", "stress", "Database that will be written to")
	insertCmd.Flags().StringVarP(&collection, "collection", "", "sensor", "Collection that will be written to")
	// insertCmd.Flags().StringVarP(&rp, "rp", "", "", "Retention Policy that will be written to")
	// insertCmd.Flags().StringVarP(&precision, "precision", "p", "n", "Resolution of data being written")
	// insertCmd.Flags().StringVarP(&consistency, "consistency", "c", "one", "Write consistency (only applicable to clusters)")
	insertCmd.Flags().IntVarP(&seriesN, "series", "s", 100000, "number of series that will be written")
	insertCmd.Flags().Uint64VarP(&pointsN, "points", "n", math.MaxUint64, "number of points that will be written")
	insertCmd.Flags().Uint64VarP(&batchSize, "batch-size", "b", 10000, "number of points in a batch")
	insertCmd.Flags().Uint64VarP(&pps, "pps", "", 200000, "Points Per Second")
	insertCmd.Flags().DurationVarP(&runtime, "runtime", "r", time.Duration(math.MaxInt64), "Total time that the test will run")
	insertCmd.Flags().DurationVarP(&tick, "tick", "", time.Second, "Amount of time between request")
	insertCmd.Flags().BoolVarP(&fast, "fast", "f", false, "Run as fast as possible")
	insertCmd.Flags().BoolVarP(&quiet, "quiet", "q", false, "Only print the write throughput")
	insertCmd.Flags().StringVar(&createCommand, "create", "", "Use a custom create database command")
	insertCmd.Flags().BoolVarP(&kapacitorMode, "kapacitor", "k", false, "Use Kapacitor mode, namely do not try to run any queries.")
	// insertCmd.Flags().IntVar(&gzip, "gzip", 0, "If non-zero, gzip write bodies with given compression level. 1=best speed, 9=best compression, -1=gzip default.")
	insertCmd.Flags().StringVar(&dump, "dump", "", "Dump to given file instead of writing over HTTP")
	insertCmd.Flags().BoolVarP(&strict, "strict", "", false, "Strict mode will exit as soon as an error or unexpected status is encountered")
	// insertCmd.Flags().BoolVarP(&tlsSkipVerify, "tls-skip-verify", "", false, "Skip verify in for TLS")
}

func client() write.Client {

	// if dump != "" {
	// 	c, err := write.NewFileClient(dump)
	// 	if err != nil {
	// 		fmt.Fprintln(os.Stderr, "Error opening file:", err)
	// 		os.Exit(1)
	// 		return c
	// 	}

	// 	return c
	// }
	return write.NewClient(host, db, collection, 10)
}
